// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/json/creator.h>
#include <nebula/ipc/json_simple.h>
#include <nebula/core/table.h>
#include <nebula/types/tensor.h>
#include <nebula/core/array.h>
#include <merak/json.h>

namespace mj = merak::json;

namespace nebula::json{


    std::shared_ptr<Array> array_from_json(const std::shared_ptr<DataType>& type,
                                         std::string_view json) {
        auto out = ipc::json::array_from_json(type, json);
        if(!out.ok()) {
            return nullptr;
        }
        return out.value_or_die();
    }

    std::shared_ptr<Array> dict_array_from_json(const std::shared_ptr<DataType>& type,
                                                       std::string_view indices_json,
                                                       std::string_view dictionary_json) {
        std::shared_ptr<Array> out;
        if(!ipc::json::dict_array_from_json(type, indices_json, dictionary_json, &out).ok()) {
            return nullptr;
        }
        return out;
    }

    std::shared_ptr<ChunkedArray> assert_chunked_array_from_json(const std::shared_ptr<DataType>& type,
                                                                 const std::vector<std::string>& json) {
        std::shared_ptr<ChunkedArray> out;
        if(!ipc::json::chunked_array_from_json(type, json, &out).ok()) {
            return nullptr;
        }
        return out;
    }


    std::shared_ptr<RecordBatch> record_batch_from_json(const std::shared_ptr<Schema>& schema,
                                                        std::string_view json) {
        // Parse as a StructArray
        auto struct_type = STRUCT(schema->fields());
        std::shared_ptr<Array> struct_array = array_from_json(struct_type, json);

        // Convert StructArray to RecordBatch
        auto rs = RecordBatch::FromStructArray(struct_array);
        if(!rs.ok()) {
            return nullptr;
        }
        return rs.value_or_die();
    }

    std::shared_ptr<Table> table_from_json(const std::shared_ptr<Schema>& schema,
                                           const std::vector<std::string>& json) {
        std::vector<std::shared_ptr<RecordBatch>> batches;
        for (const std::string& batch_json : json) {
            batches.push_back(record_batch_from_json(schema, batch_json));
        }
        auto rs = Table::from_record_batches(schema, batches);
        if(!rs.ok()) {
            return nullptr;
        }
        return rs.value_or_die();
    }

    std::shared_ptr<Scalar> scalar_from_json(const std::shared_ptr<DataType>& type,
                                                    std::string_view json) {
        std::shared_ptr<Scalar> out;
        auto r = ipc::json::scalar_from_json(type, json, &out);
        if(!r.ok()) {
            return nullptr;
        }
        return out;
    }

    std::shared_ptr<Scalar> dict_scalar_from_json(const std::shared_ptr<DataType>& type,
                                                         std::string_view index_json,
                                                         std::string_view dictionary_json) {
        std::shared_ptr<Scalar> out;
        if(!ipc::json::dict_scalar_from_json(type, index_json, dictionary_json, &out).ok()) {
            return nullptr;
        }
        return out;
    }

    std::shared_ptr<Tensor> tensor_from_json(const std::shared_ptr<DataType>& type,
                                                    std::string_view data, std::string_view shape,
                                                    std::string_view strides,
                                                    std::string_view dim_names) {
        std::shared_ptr<Array> array = array_from_json(type, data);

        mj::Document json_shape;
        json_shape.parse(shape.data(), shape.length());
        std::vector<int64_t> shape_vector;
        for (auto& x : json_shape.get_array()) {
            shape_vector.emplace_back(x.get_int64());
        }
        mj::Document json_strides;
        json_strides.parse(strides.data(), strides.length());
        std::vector<int64_t> strides_vector;
        for (auto& x : json_strides.get_array()) {
            strides_vector.emplace_back(x.get_int64());
        }
        mj::Document json_dim_names;
        json_dim_names.parse(dim_names.data(), dim_names.length());
        std::vector<std::string> dim_names_vector;
        for (auto& x : json_dim_names.get_array()) {
            dim_names_vector.emplace_back(x.get_string());
        }
        auto rs = Tensor::create(type, array->data()->buffers[1], shape_vector, strides_vector,dim_names_vector);
        if(!rs.ok()) {
            return nullptr;
        }
        return rs.value_or_die();
    }

    std::shared_ptr<Tensor> tensor_from_json(const std::shared_ptr<DataType>& type,
                                                    std::string_view data,
                                                    const std::vector<int64_t>& shape,
                                                    const std::vector<int64_t>& strides,
                                                    const std::vector<std::string>& dim_names) {
        std::shared_ptr<Array> array = array_from_json(type, data);
        auto rs = Tensor::create(type, array->data()->buffers[1], shape, strides, dim_names);
        if(!rs.ok()) {
            return nullptr;
        }
        return rs.value_or_die();
    }


}  // namespace nebula::json